# create_clf.py
import cv2
import numpy as np
from joblib import dump
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import classification_report
import logging

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)

def apply_filters(image_path):
    """对输入图像应用多个滤波器并返回特征矩阵"""
    # 读取原始图像
    image = cv2.imread(image_path)
    
    # 应用滤波器
    mean_filtered = cv2.blur(image, (5, 5))
    gaussian_filtered = cv2.GaussianBlur(image, (5, 5), 0)
    
    sobelx = cv2.Sobel(image, cv2.CV_64F, 1, 0, ksize=5)
    sobely = cv2.Sobel(image, cv2.CV_64F, 0, 1, ksize=5)
    sobel_combined = cv2.magnitude(sobelx, sobely)
    
    canny_filtered = cv2.Canny(image, 100, 200)
    
    # 转换为灰度图并重塑
    img1 = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY).reshape(1, -1)
    img2 = cv2.cvtColor(mean_filtered, cv2.COLOR_BGR2GRAY).reshape(1, -1)
    img3 = cv2.cvtColor(gaussian_filtered, cv2.COLOR_BGR2GRAY).reshape(1, -1)
    img4 = cv2.cvtColor(sobel_combined.astype(np.uint8), cv2.COLOR_BGR2GRAY).reshape(1, -1)
    img5 = canny_filtered.reshape(1, -1)
    
    return np.vstack((img1, img2, img3, img4, img5)).T

def get_labels_from_segment(segment_path):
    """从分割图像中获取标签"""
    # 读取分割图像
    segment = cv2.imread(segment_path)
    hsv_segment = cv2.cvtColor(segment, cv2.COLOR_BGR2HSV)
    
    # 定义颜色的HSV阈值
    red_lower = np.array([0, 120, 70])
    red_upper = np.array([10, 255, 255])
    green_lower = np.array([40, 40, 40])
    green_upper = np.array([80, 255, 255])
    blue_lower = np.array([110, 50, 50])
    blue_upper = np.array([130, 255, 255])
    yellow_lower = np.array([20, 100, 100])
    yellow_upper = np.array([30, 255, 255])
    
    # 创建掩膜
    red_mask = cv2.inRange(hsv_segment, red_lower, red_upper)
    green_mask = cv2.inRange(hsv_segment, green_lower, green_upper)
    blue_mask = cv2.inRange(hsv_segment, blue_lower, blue_upper)
    yellow_mask = cv2.inRange(hsv_segment, yellow_lower, yellow_upper)
    
    # 创建标签
    label = np.ones(red_mask.shape)
    label[red_mask == 255] = 1
    label[green_mask == 255] = 2
    label[blue_mask == 255] = 3
    label[yellow_mask == 255] = 4
    
    return label.reshape(-1)

def main():
    # 设置路径
    image_path = r'D:\project8\examination_data\train_data\Sandstone_1.png'
    segment_path = r'D:\project8\examination_data\train_data\Sandstone_1_segment.png'
    
    # 获取特征和标签
    logging.info('开始提取特征...')
    X = apply_filters(image_path)
    logging.info('开始获取标签...')
    y = get_labels_from_segment(segment_path)
    
    # 获取图片形状
    image = cv2.imread(image_path)
    logging.info('图片形状: %s', str(image.shape))
    
    # 数据分割
    logging.info('开始分割训练集和测试集...')
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
    
    # 特征标准化
    logging.info('特征标准化...')
    scaler = StandardScaler()
    X_train = scaler.fit_transform(X_train)
    X_test = scaler.transform(X_test)
    
    # 训练模型
    logging.info('开始训练随机森林模型...')
    clf = RandomForestClassifier(
        n_estimators=50,
        max_depth=10,
        min_samples_split=5,
        random_state=42
    )
    clf.fit(X_train, y_train)
    logging.info('完成模型训练')
    
    # 评估模型
    logging.info('评估模型...')
    y_pred = clf.predict(X_test)
    print(classification_report(y_test, y_pred))
    
    # 保存模型和scaler
    logging.info('保存模型和scaler...')
    model_data = {
        'clf': clf,
        'scaler': scaler
    }
    dump(model_data, r'D:\project8\examination_data\model.joblib', compress=3)
    logging.info('模型保存完成')

if __name__ == '__main__':
    main()